{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mask R-CNN - Human pose estimation\n",
    "\n",
    "Inspect and visualize data loading and pre-processing code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n",
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import itertools\n",
    "import math\n",
    "import logging\n",
    "import json\n",
    "import re\n",
    "import random\n",
    "from collections import OrderedDict\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import matplotlib.lines as lines\n",
    "from matplotlib.patches import Polygon\n",
    "\n",
    "import utils\n",
    "import visualize\n",
    "from visualize import display_images\n",
    "import model as modellib\n",
    "from model import log\n",
    "\n",
    "%matplotlib inline \n",
    "\n",
    "ROOT_DIR = os.getcwd()\n",
    "MODEL_DIR = os.path.join(ROOT_DIR, \"logs\")\n",
    "# Local path to trained weights file\n",
    "COCO_MODEL_PATH = os.path.join(ROOT_DIR, \"mask_rcnn_coco.h5\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configurations\n",
    "\n",
    "Run one of the code blocks below to import and load the configurations to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MS COCO Dataset\n",
    "import coco\n",
    "config = coco.CocoConfig()\n",
    "COCO_DIR = \"/home/xiaowu/Documents/Mask_RCNN_Humanpose-master/images\"  # TODO: enter your own path here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=36.34s)\n",
      "creating index...\n",
      "index created!\n",
      "Skeleton: (19, 2)\n",
      "Keypoint names: (17,)\n",
      "loading annotations into memory...\n",
      "Done (t=4.62s)\n",
      "creating index...\n",
      "index created!\n",
      "Skeleton: (19, 2)\n",
      "Keypoint names: (17,)\n",
      "Train Keypoints Image Count: 64115\n",
      "Train Keypoints Class Count: 2\n",
      "  0. BG                                                \n",
      "  1. person                                            \n",
      "Val Keypoints Image Count: 2693\n",
      "Val Keypoints Class Count: 2\n",
      "  0. BG                                                \n",
      "  1. person                                            \n"
     ]
    }
   ],
   "source": [
    "# MS COCO Dataset\n",
    "import coco\n",
    "config = coco.CocoConfig()\n",
    "COCO_DIR = \"/home/xiaowu/Documents/Mask_RCNN_Humanpose-master/images\"  # TODO: enter value here\n",
    "# Load dataset\n",
    "assert config.NAME == \"coco\"\n",
    "# Training dataset\n",
    "# load person keypoints dataset\n",
    "train_dataset_keypoints = coco.CocoDataset(task_type=\"person_keypoints\")\n",
    "train_dataset_keypoints.load_coco(COCO_DIR, \"train\")\n",
    "train_dataset_keypoints.prepare()\n",
    "\n",
    "#Validation dataset\n",
    "val_dataset_keypoints = coco.CocoDataset(task_type=\"person_keypoints\")\n",
    "val_dataset_keypoints.load_coco(COCO_DIR, \"val\")\n",
    "val_dataset_keypoints.prepare()\n",
    "\n",
    "print(\"Train Keypoints Image Count: {}\".format(len(train_dataset_keypoints.image_ids)))\n",
    "print(\"Train Keypoints Class Count: {}\".format(train_dataset_keypoints.num_classes))\n",
    "for i, info in enumerate(train_dataset_keypoints.class_info):\n",
    "    print(\"{:3}. {:50}\".format(i, info['name']))\n",
    "\n",
    "print(\"Val Keypoints Image Count: {}\".format(len(val_dataset_keypoints.image_ids)))\n",
    "print(\"Val Keypoints Class Count: {}\".format(val_dataset_keypoints.num_classes))\n",
    "for i, info in enumerate(val_dataset_keypoints.class_info):\n",
    "    print(\"{:3}. {:50}\".format(i, info['name']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ceate Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model object in inference mode.\n",
    "model = modellib.MaskRCNN(mode=\"training\", model_dir=MODEL_DIR, config=config)\n",
    "\n",
    "# Load weights trained on MS-COCO\n",
    "model.load_weights(COCO_MODEL_PATH, by_name=True,exclude=[\"mrcnn_class_logits\", \"mrcnn_bbox_fc\", \n",
    "                                \"mrcnn_bbox\", \"mrcnn_mask\"])\n",
    "print(\"Loading weights from \", model_path)\n",
    "# model.keras_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Training - Stage 1\n",
    "print(\"Train heads\")\n",
    "model.train(train_dataset_keypoints, val_dataset_keypoints,\n",
    "            learning_rate=config.LEARNING_RATE,\n",
    "            epochs=15,\n",
    "            layers='heads')\n",
    "# Training - Stage 2\n",
    "# Finetune layers from ResNet stage 4 and up\n",
    "print(\"Training Resnet layer 4+\")\n",
    "model.train(train_dataset_keypoints, val_dataset_keypoints,\n",
    "            learning_rate=config.LEARNING_RATE / 10,\n",
    "            epochs=20,\n",
    "            layers='4+')\n",
    "# Training - Stage 3\n",
    "# Finetune layers from ResNet stage 3 and up\n",
    "print(\"Training Resnet layer 3+\")\n",
    "model.train(train_dataset_keypoints, val_dataset_keypoints,\n",
    "            learning_rate=config.LEARNING_RATE / 100,\n",
    "            epochs=100,\n",
    "            layers='all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = os.path.join(MODEL_DIR, \"mobile_mask_rcnn_coco_humanpose.h5\")\n",
    "model.keras_model.save_weights(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
